#!/usr/bin/env python3
###
# @Author: Logan.Li
# @Gitee: https://gitee.com/attacker
# @email: admin@attacker.club
# @Date: 2025-03-14 08:30:30
# @LastEditTime: 2025-03-14 08:30:35
# @Description: 统一告警通知模块
###

import json
import logging
import requests
from abc import ABC, abstractmethod
from typing import Dict, Any

class NotifierBase(ABC):
    """通知器基类"""
    
    @abstractmethod
    def send(self, title: str, content: str) -> bool:
        """发送通知"""
        pass

class DingTalkNotifier(NotifierBase):
    """钉钉群通知"""
    
    def __init__(self, webhook: str, secret: str = None):
        self.webhook = webhook
        self.secret = secret
        
    def send(self, title: str, content: str) -> bool:
        try:
            headers = {'Content-Type': 'application/json'}
            data = {
                "msgtype": "markdown",
                "markdown": {
                    "title": title,
                    "text": f"### {title}\n{content}"
                }
            }
            response = requests.post(self.webhook, headers=headers, json=data)
            return response.status_code == 200
        except Exception as e:
            logging.error(f"钉钉通知发送失败: {e}")
            return False

class LarkNotifier(NotifierBase):
    """飞书通知"""
    
    def __init__(self, webhook: str):
        self.webhook = webhook
        
    def send(self, title: str, content: str) -> bool:
        try:
            headers = {'Content-Type': 'application/json'}
            data = {
                "msg_type": "post",
                "content": {
                    "post": {
                        "zh_cn": {
                            "title": title,
                            "content": [[{"tag": "text", "text": content}]]
                        }
                    }
                }
            }
            response = requests.post(self.webhook, headers=headers, json=data)
            return response.status_code == 200
        except Exception as e:
            logging.error(f"飞书通知发送失败: {e}")
            return False

class TelegramNotifier(NotifierBase):
    """Telegram通知"""
    
    def __init__(self, bot_token: str, chat_id: str):
        self.bot_token = bot_token
        self.chat_id = chat_id
        self.api_url = f"https://api.telegram.org/bot{bot_token}/sendMessage"
        
    def send(self, title: str, content: str) -> bool:
        try:
            data = {
                "chat_id": self.chat_id,
                "text": f"*{title}*\n{content}",
                "parse_mode": "Markdown"
            }
            response = requests.post(self.api_url, json=data)
            return response.status_code == 200
        except Exception as e:
            logging.error(f"Telegram通知发送失败: {e}")
            return False

class NotificationManager:
    """通知管理器"""
    
    def __init__(self, config: Dict[str, Any]):
        self.notifiers = []
        self._init_notifiers(config)
        
    def _init_notifiers(self, config: Dict[str, Any]):
        """初始化所有配置的通知器"""
        if config.get('dingtalk', {}).get('enabled'):
            self.notifiers.append(
                DingTalkNotifier(
                    config['dingtalk']['webhook'],
                    config['dingtalk'].get('secret')
                )
            )
            
        if config.get('lark', {}).get('enabled'):
            self.notifiers.append(
                LarkNotifier(config['lark']['webhook'])
            )
            
        if config.get('telegram', {}).get('enabled'):
            self.notifiers.append(
                TelegramNotifier(
                    config['telegram']['bot_token'],
                    config['telegram']['chat_id']
                )
            )
            
    def send_notification(self, title: str, content: str) -> bool:
        """发送通知到所有配置的通知渠道"""
        success = True
        for notifier in self.notifiers:
            if not notifier.send(title, content):
                success = False
        return success

# 使用示例
if __name__ == "__main__":
    # 配置示例
    config = {
        "dingtalk": {
            "enabled": True,
            "webhook": "https://oapi.dingtalk.com/robot/send?access_token=xxx",
            "secret": "your_secret"
        },
        "lark": {
            "enabled": True,
            "webhook": "https://open.feishu.cn/open-apis/bot/v2/hook/xxx"
        },
        "telegram": {
            "enabled": True,
            "bot_token": "your_bot_token",
            "chat_id": "your_chat_id"
        }
    }
    
    # 创建通知管理器
    notification_manager = NotificationManager(config)
    
    # 发送测试通知
    notification_manager.send_notification(
        "监控告警测试",
        "这是一条测试告警消息\n- CPU使用率: 90%\n- 内存使用率: 85%"
    )
